from pymongo import MongoClient

from settings import MONGO_URL
from utils.log import logger
from domain import Proxy

"""
实现对数据库增删改查相关操作
"""


class MongoPool(object):
    def __init__(self):
        self.client = MongoClient(MONGO_URL)
        self.proxies = self.client['app']['proxies']

    def __del__(self):
        self.client.close()

    def insert_one(self, proxy):
        count = self.proxies.count_documents({'_id': proxy.ip})
        if count == 0:
            dic = proxy.__dict__
            dic['_id'] = proxy.ip
            self.proxies.insert_one(dic)
            logger.info('插入了代理：{}'.format(proxy))
        else:
            logger.warning('已存在代理：{}'.format(proxy))

    def update_one(self, proxy):
        if proxy:
            self.proxies.update_one({'_id': proxy.ip}, {"$set": proxy.__dict__})
        else:
            logger.error('需要更新的代理对象不正确')

    def delete_one(self, proxy):
        if proxy:
            self.proxies.delete_one({'-id': proxy.ip})
        else:
            logger.error('需要删除的代理对象不正确')

    def find_all(self):
        cursor = self.proxies.find()
        for item in cursor:
            item.pop('_id')

            proxy = Proxy(**item)
            yield proxy

    def select_one(self, id):
        first = self.proxies.find_one(filter={'_id': id})
        return first


mongo = MongoPool()

if __name__ == '__main__':
    # proxy = Proxy('182.84.81.128', '9999')
    # proxy = Proxy('114.231.42.126', '8888')

    # mongo.insert_one(proxy)

    # for proxy in mongo.find_all():
    #     print(proxy)

    proxy = mongo.select_one('123.205.24.245')
    if proxy:
        print(proxy['ip'])
    else:
        print("test")
